/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "model.h"

#include <errno.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

#define MODEL_VERSION 1ULL

typedef struct _Header Header;
struct _Header 
{
  uint64_t      version;
  uint64_t      n_items;
  uint32_t      n_labels;
  uint32_t      n_workers;
  uint32_t      symmetric;
  uint32_t      pad;
};

typedef struct _ModelImpl ModelImpl;
struct _ModelImpl 
{
  int                                   fd;
  uint8_t*                              base;
  ParameterizedNormalDistribution       prioralpha;
  NormalDistribution                    priorlogbeta;
  NormalDistribution                    priorgamma;
};

static float*
alpha_pointer (uint8_t* base)
{
  Header* h = (Header*) base;

  return (float*) (
         base 
       + sizeof (*h) 
       + (h->n_labels * sizeof (float))                             /* priorz */
       + (h->n_labels * h->n_labels * sizeof (float))               /* gamma */
    );
}

static float*
gamma_pointer (uint8_t* base)
{
  Header* h = (Header*) base;
  return (float*) (
         base 
       + sizeof (*h) 
       + (h->n_labels * sizeof (float))                             /* priorz */
    );
}

static const float*
priorz_pointer (const uint8_t* base)
{
  const Header* h = (const Header*) base;
  return (float*) (
         base
       + sizeof (*h) 
    );
}

static size_t
file_size (const Header* h)
{
  return sizeof (*h)                                                /* header */
       + (h->n_labels * sizeof (float))                             /* priorz */
       + (h->n_labels * h->n_labels * sizeof (float))               /* gamma */
       + (h->n_workers * h->n_labels * h->n_labels * sizeof (float))/* alpha */
    ;
}

Model*
model_open (const char* pathname,
            bool        read_only)
{
  Header tmp_hdr;
  int fd = open (pathname, (read_only) ? O_RDONLY : O_RDWR);

  if (fd < 0)
    {
      fprintf (stderr, "can't open '%s': %s\n", pathname, strerror (errno));
      return NULL;
    }

  if (lseek (fd, 0, SEEK_SET) == (off_t) -1 ||
      read (fd, &tmp_hdr, sizeof (tmp_hdr)) < (ssize_t) sizeof (tmp_hdr))
    {
      fprintf (stderr,
               "error reading header of model '%s': %s\n",
               pathname,
               strerror (errno));
      goto FAIL;
    }

  if (tmp_hdr.version != MODEL_VERSION)
    {
      fprintf (stderr, 
               "model version mismatch (want %llu got %llu)\n",
               MODEL_VERSION,
               (long long unsigned int) tmp_hdr.version);
      goto FAIL;
    }

  uint8_t* base = mmap (NULL,
                        file_size (&tmp_hdr),
                        PROT_READ | (read_only ? 0 : PROT_WRITE),
                        MAP_SHARED,
                        fd,
                        0);

  if (base == (uint8_t*) (-1))
    {
      fprintf (stderr,
               "can't mmap %zu bytes of '%s': %s\n",
               file_size (&tmp_hdr),
               pathname,
               strerror (errno));
      goto FAIL;
    }

  ModelImpl* impl = (ModelImpl*) malloc (sizeof (ModelImpl));
  if (impl == NULL)
    {
      fprintf (stderr,
               "wtf, can't alloc %zu bytes (?): %s\n",
               sizeof (ModelImpl),
               strerror (errno));
      goto FAIL;
    }
  impl->fd = fd;
  impl->base = base;
  impl->prioralpha = parameterized_normal_distribution (1);
  impl->priorlogbeta = normal_distribution (0, 1);
  impl->priorgamma = normal_distribution (0, 1);

  Model* model = (Model*) malloc (sizeof (Model));
  if (model == NULL)
    {
      fprintf (stderr,
               "wtf, can't alloc %zu bytes (?): %s\n",
               sizeof (Model),
               strerror (errno));
      free (impl);
      goto FAIL;
    }

  model->n_items = ((Header*) base)->n_items;
  model->n_labels = ((Header*) base)->n_labels;
  model->n_workers = ((Header*) base)->n_workers;
  model->alpha = alpha_pointer (base);
  model->gamma = gamma_pointer (base);
  model->logpriorz = priorz_pointer (base);
  model->prioralpha = &impl->prioralpha;
  model->priorlogbeta = &impl->priorlogbeta.base;
  model->priorgamma = &impl->priorgamma.base;
  model->symmetric = ((Header*) base)->symmetric;
  model->impl = impl;

  return model;

FAIL:
  close (fd);
  return NULL;
}

Model*
model_creat (const char*        pathname,
             unsigned int       n_items,
             unsigned int       n_labels,
             unsigned int       n_workers,
             const float*       logpriorz,
             bool               symmetric)
{
  Header header = { .version = MODEL_VERSION,
                    .n_items = n_items,
                    .n_labels = n_labels,
                    .n_workers = n_workers,
                    .symmetric = symmetric,
                    .pad = 0 };
  int fd = open (pathname,
                 O_RDWR | O_CREAT | O_EXCL,
                 S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH);

  if (fd < 0)
    {
      fprintf (stderr, "can't creat '%s': %s\n", pathname, strerror (errno));
      return NULL;
    }

  off_t end = lseek (fd, 0, SEEK_END);

  if (end != 0 ||
      ftruncate (fd, file_size (&header)) < 0 ||
      lseek (fd, 0, SEEK_SET) == (off_t) -1 ||
      write (fd, &header, sizeof (header)) < (ssize_t) sizeof (header) ||
      write (fd, logpriorz, n_labels * sizeof (float)) < (ssize_t) (n_labels * sizeof (float)))
    {
      fprintf (stderr,
               "error writing header of model '%s': %s\n",
               pathname,
               strerror (errno));
      unlink (pathname);
      close (fd);
      return NULL;
    }

  close (fd);
  return model_open (pathname, false);
}

void
model_close (Model* model)
{
  ModelImpl* impl = model->impl;

  munmap (impl->base, file_size ((Header*) impl->base));
  close (impl->fd);
  free (impl);
  free (model);
}
